import PIL.Image as Image
import gradio as gr
import matplotlib
matplotlib.use('TkAgg')
from ultralytics import YOLO
from obb_predict import get_point_info,get_theho_result

# Load a model
model = YOLO("./best.pt")  # pretrained YOLO11n model


def predict_image(img,max_value,zero_offset):
    results = model.predict(
        source=img,
        show_labels=True,
        show_conf=True,
        imgsz=640,
        conf=0.01
    )
    #     obb = result.obb  # Oriented boxes object for OBB outputs
    #
    #     info = get_point_info(obb)

    results2 = model.predict(
        source=img,
        show_labels=True,
        show_conf=True,
        imgsz=640,
    )

    predict_value = -1
    for r,r2 in zip(results,results2):
        im_array = r2.plot()
        im = Image.fromarray(im_array[..., ::-1])
        obb = r.obb  # Oriented boxes object for OBB outputs
        info = get_point_info(obb)
        predict_value = get_theho_result(info,float(max_value),float(zero_offset))

    return im,predict_value

interface = gr.Interface(
    fn=predict_image,
    inputs=[
        gr.Image(type="pil", label="上传照片"),
        gr.Textbox(value="1.6",placeholder="请输入压力表最大刻度",label="最大刻度"),
        gr.Textbox(value="5",placeholder="请输入偏移量",label="0刻度偏移量")
    ],
    outputs=[gr.Image(type="pil", label="标记结果"),gr.Text(label="预测结果")],
    title="压力表刻度检测",
    description='''压力表刻度检测''',
    examples=[
        ["./samples/1.jpg","60","1"],
        ["./samples/2.jpg", "1.6","0.025"],
        ["./samples/3.jpg", "60","1"],
        ["./samples/4.jpg", "1.6","0.025"],
        ["./samples/5.jpg", "25","0.5"],

    ]
)

if __name__ == "__main__":
    interface.launch()
